# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import math


def adjust_learning_rate(optimizer, epoch, args):
    """
    调整学习率，采用预热（warmup）后半周期余弦衰减策略。
    在预热阶段，学习率线性增加；之后，使用半周期余弦衰减策略。

    参数：
        optimizer (torch.optim.Optimizer): 用于训练的优化器。
        epoch (int): 当前的训练轮数（epoch）。
        args (Namespace): 存储训练配置的命名空间对象，包含学习率、预热轮数、总轮数等信息。

    返回：
        float: 调整后的学习率。
    """
    # 如果当前epoch在预热期内，则学习率线性增加
    if epoch < args.warmup_epochs:
        lr = args.lr * epoch / args.warmup_epochs  # 线性增加学习率
    else:
        # 超过预热期后，使用余弦衰减策略
        lr = args.min_lr + (args.lr - args.min_lr) * 0.5 * \
             (1. + math.cos(math.pi * (epoch - args.warmup_epochs) / (args.epochs - args.warmup_epochs)))
        # 其中，args.lr是初始学习率，args.min_lr是最小学习率，args.epochs是总训练轮数

    # 遍历优化器的参数组，更新每个参数组的学习率
    for param_group in optimizer.param_groups:
        # 如果当前参数组有"lr_scale"，则学习率将乘以该比例
        if "lr_scale" in param_group:
            param_group["lr"] = lr * param_group["lr_scale"]
        else:
            # 否则直接设置为计算后的学习率
            param_group["lr"] = lr

    # 返回调整后的学习率
    return lr

